﻿#include "pch.h"
#include <windows.h>
#include <winioctl.h>
#include <setupapi.h>
#include <devguid.h>
#include "DiskManager.h"

#pragma comment(lib, "setupapi.lib")

DiskManager::DiskManager()
{}

DiskManager::~DiskManager()
{}

BOOL DiskManager::Enum()
{
    for(int driveNumber = 0; driveNumber < 26; ++driveNumber) {
        BOOL bRet = TRUE;
        char drivePath[MAX_PATH] = { 0 };
        sprintf_s(drivePath, "\\\\.\\PHYSICALDRIVE%d", driveNumber);

        HANDLE hDevice = CreateFileA(drivePath, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
        if(hDevice == INVALID_HANDLE_VALUE)
            return FALSE;

        STORAGE_PROPERTY_QUERY storagePropertyQuery;
        ZeroMemory(&storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY));
        storagePropertyQuery.PropertyId = StorageDeviceProperty;
        storagePropertyQuery.QueryType = PropertyStandardQuery;

        STORAGE_DESCRIPTOR_HEADER storageDescriptorHeader = { 0 };
        DISK_GEOMETRY diskGeometry;
        DWORD bytesReturned = 0;

        // 获取磁盘属性信息
        if(DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY),
            &storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER), &bytesReturned, NULL) == FALSE)
            return FALSE;

        if(storageDescriptorHeader.Size > 0) {
            char* buffer = new char[storageDescriptorHeader.Size];
            if(DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY,
                &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY),
                buffer, storageDescriptorHeader.Size, &bytesReturned, NULL)) {
                STORAGE_DEVICE_DESCRIPTOR* deviceDescriptor = reinterpret_cast<STORAGE_DEVICE_DESCRIPTOR*>(buffer);

                // 可移动磁盘
                if(deviceDescriptor->RemovableMedia) {

                    // 磁盘序号
                    CString szDriveNumber;
                    szDriveNumber.Format(_T("Disk:%d"), driveNumber);

                    // 磁盘容量
                    LONGLONG nTotalBytes = 0;
                    LONGLONG nBlockCount = 0;
                    if(DeviceIoControl(hDevice, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, &diskGeometry, sizeof(diskGeometry), &bytesReturned, NULL)) {
                        nBlockCount = diskGeometry.Cylinders.QuadPart *
                                      diskGeometry.TracksPerCylinder *
                                      diskGeometry.SectorsPerTrack;

                        nTotalBytes = diskGeometry.BytesPerSector * nBlockCount;
                    }

                    DISK_DRIVE_INFO info;
                    info.szDriveNumber = szDriveNumber;
                    info.szDriveFilePath = drivePath;
                    info.nTotalBytes = nTotalBytes;
                    info.nBlockCount = nBlockCount;
                    info.szDriveName =
                        CString(buffer + deviceDescriptor->VendorIdOffset) +
                        CString(buffer + deviceDescriptor->ProductIdOffset);
                    m_diskDrives.AddTail(info);
                }
            }
            delete[] buffer;
        }
        CloseHandle(hDevice);
    }
    return TRUE;
}

BOOL DiskManager::GetInfo(LPDISK_DRIVE_INFO lpDiskDriveInfo, UINT nIndex) const
{
    if(nIndex < 0 && nIndex > (UINT)m_diskDrives.GetSize() - 1)
        return FALSE;
   
    POSITION pos;
    if((pos = m_diskDrives.FindIndex(nIndex)) != NULL) {
        *lpDiskDriveInfo = m_diskDrives.GetAt(pos);
        return TRUE;
    }
    return FALSE;
}

UINT DiskManager::GetCount() const
{
    return m_diskDrives.GetSize();
}

BOOL DiskManager::InitDisk(const CString& pszDiskPath) const
{
    BOOL bRet = TRUE;
    static CREATE_DISK raw;

    HANDLE hDevice = CreateFile(
        (LPCTSTR)pszDiskPath,
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        0,
        OPEN_EXISTING,
        FILE_FLAG_NO_BUFFERING,
        NULL
    );

    if(hDevice == INVALID_HANDLE_VALUE)
        return FALSE;

    DWORD nBytesReturned;
    if(!DeviceIoControl(hDevice, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &nBytesReturned, NULL)) {
        bRet = FALSE;
        goto exit;
    }

    if(!DeviceIoControl(hDevice, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &nBytesReturned, NULL)) {
        bRet = FALSE;
        goto exit;
    }

    raw.PartitionStyle = PARTITION_STYLE_GPT;
    if(!DeviceIoControl(hDevice, IOCTL_DISK_CREATE_DISK, &raw, sizeof(raw), NULL, 0, &nBytesReturned, NULL)) {
        bRet = FALSE;
        goto exit;
    }

    //if(!DeviceIoControl(hDevice, IOCTL_DISK_DELETE_DRIVE_LAYOUT, NULL, 0, NULL, 0, NULL, NULL)) {
    //    bRet = FALSE;
    //    goto exit;
    //}

    if(!DeviceIoControl(hDevice, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &nBytesReturned, NULL)) {
        bRet = FALSE;
        goto exit;
    }

    //if(!DeviceIoControl(hDevice, IOCTL_STORAGE_EJECT_MEDIA, NULL, 0, NULL, 0, &nBytesReturned, NULL)) {
    //    bRet = FALSE;
    //    goto exit;
    //}

    if(!DeviceIoControl(hDevice, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &nBytesReturned, NULL)) {
        bRet = FALSE;
        goto exit;
    }

exit:
    CloseHandle(hDevice);
    return bRet;
}